package com.bsoft.advance.core.register.zk;

import com.bsoft.advance.core.balancer.LoadBalance;
import com.bsoft.advance.core.common.ServiceInfo;
import com.bsoft.advance.core.register.DiscoveryService;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import org.springframework.util.CollectionUtils;

import java.util.Collection;
import java.util.stream.Collectors;

/**
 * zk服务发现
 */
@Slf4j
public class ZookeeperDiscoveryService implements DiscoveryService {

    private ServiceDiscovery<ServiceInfo> serviceDiscovery;

    private LoadBalance loadBalance;

    public ZookeeperDiscoveryService(String registryAddr, LoadBalance loadBalance) {
        this.loadBalance = loadBalance;
        try {
            // 获取zk客户端
            CuratorFramework client = CuratorUtil.getZkClient(registryAddr);
            // json序列化
            JsonInstanceSerializer<ServiceInfo> serializer = new JsonInstanceSerializer<>(ServiceInfo.class);
            this.serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceInfo.class)
                    .client(client)
                    .serializer(serializer)
                    .basePath(CuratorUtil.ZK_REGISTER_ROOT_PATH)
                    .build();
            this.serviceDiscovery.start();
        } catch (Exception e) {
            log.error("serviceDiscovery start error :{}", e);
        }
    }

    /**
     * 服务发现
     */
    @Override
    public ServiceInfo discovery(String serviceName) throws Exception {
        // 根据服务名，找出所有实例
        Collection<ServiceInstance<ServiceInfo>> serviceInstances = serviceDiscovery.queryForInstances(serviceName);
        // 根据负载均衡策略，选择一个
        return CollectionUtils.isEmpty(serviceInstances) ? null : loadBalance.doSelect(serviceInstances.stream().map(ServiceInstance::getPayload).collect(Collectors.toList()));
    }

}
